import copy

import numpy as np
import pydicom
from PIL import Image
from matplotlib import pyplot

from models.SENet.SEUnet import SimSEUnet
from models.unet.simUnet import simUNet
from utils.train.imageUtils import maybe_rotate, letterbox_image, normalize
from utils.evaluation.prediction import Prediction

if __name__ == "__main__":

    num_classes = 'both'
    assert num_classes in ['inner', 'outer', 'both']

    model_path = "../RVSC/checkpoint_unet_1/Epoch66-mean_dsc0.8710.pth"
    model_image_size = (256, 256, 1)
    cuda = True
    blend = True

    if num_classes == 'both':
        num_classes = 3
    else:
        num_classes = 2
    net = simUNet(in_channels=model_image_size[-1],
                  out_channels=model_image_size[-1],
                  init_features=64).eval()
    unetPrediction = Prediction(net=net,
                                model_path=model_path,
                                model_image_size=model_image_size,
                                num_classes=num_classes,
                                cuda=cuda,
                                blend=blend)

    while True:
        img = input('Input image filename:')
        try:
            image = pydicom.read_file(img)
            image = maybe_rotate(image.pixel_array)
            image = Image.fromarray(image.astype(np.uint8)).convert('L')

            old_img = copy.deepcopy(image)
            image, nw, nh = letterbox_image(image=image,
                                            size=(model_image_size[1], model_image_size[0]),
                                            image_type='L')

            pyplot.imshow(image)
            pyplot.show()

            images = normalize([[np.array(image)]], axis=(2, 3))
        except FileNotFoundError:
            print('Open Error! Try again!')
            continue
        else:
            r_image = unetPrediction.detect_image2(old_img=old_img, images=images, nw=nw, nh=nh)
            r_image.show()
